import shutil
import cv2
from flask import Flask, jsonify, make_response, render_template, Response, request, send_file
import time
import os
import Raster
import logging
log = logging.getLogger('werkzeug')
log.setLevel(logging.ERROR)
app = Flask(__name__)

img_idx = 0
num = 0

@app.route('/')
def index():
    return render_template('index.html', number=0)

@app.route('/images/background.png')
def background():
    return send_file('templates/static/images/background.png', mimetype='image/png')

@app.route('/images/chair.png')
def chair():
    return send_file('templates/static/images/chair.png', mimetype='image/png')

@app.route('/images/hotdog.png')
def hotdog():
    return send_file('templates/static/images/hotdog.png', mimetype='image/png')

@app.route('/images/drum.png')
def drum():
    return send_file('templates/static/images/drum.png', mimetype='image/png')

@app.route('/images/lego.png')
def lego():
    return send_file('templates/static/images/lego.png', mimetype='image/png')

@app.route('/images/ship.png')
def ship():
   return send_file('templates/static/images/ship.png', mimetype='image/png')

@app.route('/images/mic.png')
def mic():
   return send_file('templates/static/images/mic.png', mimetype='image/png')

clear_img_flag = False
def gen():
    global img_idx
    global num
    global clear_img_flag
    while True:
        while not clear_img_flag:
            img_path = 'output/' + str(img_idx).zfill(5) + '.png'
            num = len(os.listdir('./output'))
            if img_idx < num - 1:
                # img_path = 'image.png'
                frame = cv2.imread(img_path)
                image = cv2.imencode('.jpg', frame)[1].tobytes()
                time.sleep(0.03)
                img_idx += 1
                img_idx %= 200
                yield (b'--frame\r\n'
                    b'Content-Type: image/jpeg\r\n\r\n' + image + b'\r\n')
            if img_idx >= 199:
                img_idx = 0

@app.route('/video_feed')
def video_feed():
    return Response(gen(), mimetype='multipart/x-mixed-replace; boundary=frame')

def generate_progress():
    i = int(num)
    time.sleep(0.1) 
    return f"data:{i}\n\n"

@app.route('/progress')
def progress():
    return app.response_class(generate_progress(), mimetype='text/event-stream; boundary=frame')

def RemoveDir(filepath):
    if not os.path.exists(filepath):
        os.mkdir(filepath)
    else:
        shutil.rmtree(filepath)
        os.mkdir(filepath)

# 获取数据
@app.route('/refresh', methods=['POST'])
def refresh():
    global clear_img_flag
    global img_idx
    clear_img_flag = True
    RemoveDir('./output/')
    # os.system("rm output/*")
    img_idx = 0
    clear_img_flag = False
    return f'删除成功'

@app.route('/changeModel', methods=['POST'])
def changeModel():
    data = request.json
    model_name = 'model/' + data['key'] + '.npz'
    # Raster.main(model_name=model_name)
    Raster.eval(data['key'])
    return jsonify({'message': data['key']})

if __name__ == '__main__':
    app.run()
    print("Hello")
